from io import BytesIO
from flask import Flask, request, send_file,send_from_directory
from utils.index import remove_background
import time
import os
app = Flask(__name__)



@app.before_request
def before_request():
    # 记录请求开始时间
    request.start_time = time.time()

@app.after_request
def after_request(response):
    # 记录请求结束时间
    if hasattr(request, 'start_time'):
        total_time = time.time() - request.start_time
        app.logger.debug(f"Request to {request.path} took {total_time:.4f} seconds")
    return response

@app.route('/remove-background', methods=['POST'])
def remove_background_api():
    # 获取上传的文件
    image_file = request.files['image']
    
    # 使用 BytesIO 创建输入流和输出流
    input_stream = BytesIO(image_file.read())
    output_stream = BytesIO()
    
    # 调用 remove_background 函数
    remove_background(input_stream, output_stream)
    
    # 返回处理后的图像
    output_stream.seek(0)  # 重置文件指针，准备发送文件流
    return send_file(output_stream, mimetype='image/png', as_attachment=True, download_name='no_bg_image.png')

@app.route('/', methods=['GET'])
def index():
    # 返回当前路径下的 index.html 文件
    return send_from_directory(os.getcwd(), 'index.html')

if __name__ == "__main__":
    app.run(debug=True,port=5000)
